import sys
sys.path.append("")
from model import clf_model
from utils.data_builder import data_builder
import numpy as np

# from matplotlib import pyplot as plt
# import seaborn as sns


# from sklearn.metrics import classification_report
# from sklearn.metrics import roc_auc_score
# from sklearn.metrics import confusion_matrix

from memory_profiler import profile
import time


EPOCHES = 10
X_PATH = "data/sentence_codes_4096_dm0.npy" # 预训练模型
ORIG_PATH = "data/train.csv"
LABLE_NAME = ["", "", "toxic","severe_toxic","obscene","threat","insult","identity_hate"]


@profile(precision=4, stream=open("logs/{}.txt".format(time.strftime("03-%Y-%m-%d %H %M %S")), 'w'))
def main():
    # 2 toxic
    # 3 severe_toxic
    # 4 obscene
    # 5 threat
    # 6 insult
    # 7 identity_hate

    builder = data_builder(X_PATH, ORIG_PATH)
    # output = None
    # 第一阶段，训练多个模型
    for col in range(2, 8):
        train_X, train_y, test_X, test_y = builder.build_one(col)

        print('train set length:\t', len(train_y))
        print('test set length:\t', len(test_y))

        clf = clf_model(epoches=EPOCHES, verbose=1, validation=(test_X, test_y))
        clf.fit(train_X, train_y)

        
        # 输出各个模型
        import pickle
        fp = open("out/{}.pickle".format(LABLE_NAME[col]), "wb")
        pickle.dump(clf, fp)
        fp.close()
    

if __name__ == "__main__":
    main()